package tree;

public class AVLTree {

    static class AVLTreeNode extends Node {

        int heightDif;

        AVLTreeNode parent;

        public AVLTreeNode(int value) {
            super(value);
            this.heightDif = 0;
            this.parent = null;
        }

        /**
         * 以前序顺序打印该节点及其子树的内容
         * @param node  节点
         */
        public static void prePrintTree(AVLTreeNode node) {
            if (node == null) {
                return;
            }
            System.out.println(node.value);
            prePrintTree((AVLTreeNode) node.left);
            prePrintTree((AVLTreeNode) node.right);
        }

        /**
         * 以中序顺序打印该节点及其子树的内容
         * @param node  节点
         */
        public static void midPrintTree(AVLTreeNode node) {
            if (node == null) {
                return;
            }
            midPrintTree((AVLTreeNode) node.left);
            System.out.println(node.value);
            midPrintTree((AVLTreeNode) node.right);
        }

        /**
         * 为某个节点计算并设置高度差
         * @param node  节点
         */
        private static void setHeightDif(AVLTreeNode node) {
            node.heightDif = getDepth(node.left) - getDepth(node.right);
        }

        /**
         * 得到某个节点的高度
         * @param node  节点
         * @return  节点的高度
         */
        private static int getDepth(Node node) {
            if (node == null) {
                return 0;
            }
            return Math.max(getDepth(node.left), getDepth(node.right)) + 1;
        }

        /**
         * 节点左旋
         * @param node  节点
         */
        private static void leftRound(AVLTreeNode node) {
            AVLTreeNode tmp = (AVLTreeNode) node.right;
            node.right = tmp.left;
            tmp.left = node;
            tmp.parent = node.parent;
            node.parent = tmp;
            if (node.right != null) {
                ((AVLTreeNode) node.right).parent = node;
            }
            if (tmp.parent != null) {
                if (tmp.parent.left == node) {
                    tmp.parent.left = tmp;
                } else {
                    tmp.parent.right = tmp;
                }
            }
            // 重新计算高度差
            setHeightDif(tmp);
            setHeightDif(node);
        }

        /**
         * 节点右旋
         * @param node  节点
         */
        private static void rightRound(AVLTreeNode node) {
            AVLTreeNode tmp = (AVLTreeNode) node.left;
            node.left = tmp.right;
            tmp.right = node;
            tmp.parent = node.parent;
            node.parent = tmp;
            if (node.left != null) {
                ((AVLTreeNode) node.left).parent = node;
            }
            if (tmp.parent != null) {
                if (tmp.parent.left == node) {
                    tmp.parent.left = tmp;
                } else {
                    tmp.parent.right = tmp;
                }
            }

            // 重新计算高度差
            setHeightDif(tmp);
            setHeightDif(node);
        }
    }

    AVLTreeNode root;

    /**
     * 插入元素
     * @param num   元素值
     */
    public void insert(int num) {
        if (root == null) {
            root = new AVLTreeNode(num);
            return;
        }

        // 找到要插入的位置的父节点
        AVLTreeNode p = root;
        boolean flag = false;
        while (!flag) {
            switch (Integer.compare(p.value, num)) {
                case 1: {
                    if (p.left != null) {
                        p = (AVLTreeNode) p.left;
                    } else {
                        flag = true;
                    }
                    break;
                }
                case -1: {
                    if (p.right != null) {
                        p = (AVLTreeNode) p.right;
                    } else {
                        flag = true;
                    }
                    break;
                }
                case 0: {
                    // 找到重复元素直接返回
                    return;
                }
            }
        }

        // 插入节点
        AVLTreeNode newNode = new AVLTreeNode(num);
        newNode.parent = p;
        if (num > p.value) {
            p.right = newNode;
        } else {
            p.left = newNode;
        }

        // 修改高度差
        p = newNode;
        flag = false;
        while (p.parent != null && !flag) {
            if (p.parent.left == p) {
                p.parent.heightDif++;
            } else {
                p.parent.heightDif--;
            }
            switch (Math.abs(p.parent.heightDif)) {
                // 从其父节点起，插入元素不会影响到更往上节点的高度差
                case 0: {
                    flag = true;
                    break;
                }
                // 原先父节点高度差为0，插入元素影响了高度差，继续向上追溯
                case 1: {
                    p = p.parent;
                    break;
                }
                // 插入破坏了平衡，需要调整平衡，调整后其父节点的高度差为0，跳出循环
                case 2: {
                    balance(p.parent, newNode);
                    flag = true;
                    break;
                }
                default: {
                    break;
                }
            }
        }
    }

    /**
     * 删除节点
     * @param num   要删除的值
     * @return  删除成功返回true，反之返回false
     */
    public boolean remove(int num) {
        AVLTreeNode p = getNode(num);
        if (p == null) {
            return false;
        }
        removeNode(p);
        return true;
    }

    /**
     * 删除节点
     * @param p 删除的节点
     */
    private void removeNode(AVLTreeNode p) {
        // 回溯的起点
        AVLTreeNode tmp = null;
        if (p.right != null && p.left != null) {
            // 找到左子树的最右节点并把值赋给p
            AVLTreeNode q = (AVLTreeNode) p.left;
            while (q.right != null) {
                q = (AVLTreeNode) q.right;
            }
            p.value = q.value;

            // 删除节点
            if (q.parent == p) {
                // 如果q的父节点就是p节点
                p.left = q.left;
                if (q.left != null) {
                    ((AVLTreeNode) q.left).parent = p;
                }
                tmp = p;
            } else {
                q.parent.right = q.left;
                if (q.left != null) {
                    ((AVLTreeNode) q.left).parent = q.parent;
                }
                tmp = q.parent;
            }
        } else if (p.left != null) {
            // 只有左子树直接将左儿子顶替自己
            if (p.parent != null) {
                if (p.parent.left == p) {
                    p.parent.left = p.left;
                } else {
                    p.parent.right = p.left;
                }
                ((AVLTreeNode) p.left).parent = p.parent;
                tmp = p.parent;
            } else {
                // 如果自己是根节点则重置根节点
                this.root = (AVLTreeNode) p.left;
            }
        } else if (p.right != null) {
            // 只有右子树直接将右儿子顶替自己
            if (p.parent != null) {
                if (p.parent.left == p) {
                    p.parent.left = p.right;
                } else {
                    p.parent.right = p.right;
                }
                ((AVLTreeNode) p.right).parent = p.parent;
                tmp = p.parent;
            } else {
                // 重置根节点
                this.root = (AVLTreeNode) p.right;
            }
        } else {
            // 左右子树都为空则清理自己的父亲节点
            if (p.parent != null) {
                if (p.parent.left == p) {
                    p.parent.left = null;
                } else {
                    p.parent.right = null;
                }
            } else {
                // 如果自己就是根节点则将根节点置为空
                this.root = null;
            }
        }

        // 回溯
        while (tmp != null) {
            int origin = tmp.heightDif;
            AVLTreeNode.setHeightDif(tmp);
            if (tmp.heightDif == origin) {
                break;
            }
            // 发生不平衡
            if (Math.abs(tmp.heightDif) == 2) {
                balance(tmp);
            }
            tmp = tmp.parent;
        }
    }

    /**
     * 根据查询值对应的节点
     * @param num 值
     * @return  对应的节点
     */
    private AVLTreeNode getNode(int num) {
        AVLTreeNode p = root;
        while (p != null) {
            switch (Integer.compare(p.value, num)) {
                case 1: {
                    p = (AVLTreeNode) p.left;
                    break;
                }
                case -1: {
                    p = (AVLTreeNode) p.right;
                    break;
                }
                case 0: {
                    return p;
                }
            }
        }
        return null;
    }

    /**
     * 是否包含某元素
     * @param num   元素
     * @return  包含的话返回true，反之返回false
     */
    public boolean contains(int num) {
        return getNode(num) != null;
    }

    /**
     * 删除节点后保持平衡
     * @param node  不平衡节点
     */
    private void balance(AVLTreeNode node) {
        if (node.heightDif == 2) {
            // 左子树高了，模拟为左儿子插入
            if (((AVLTreeNode) node.left).heightDif > 0) {
                // 模拟插入的是左儿子的左子树
                AVLTreeNode.rightRound(node);
            } else {
                // 模拟插入的是左儿子的右子树
                AVLTreeNode.leftRound(((AVLTreeNode) node.left));
                AVLTreeNode.rightRound(node);
            }
        } else {
            // 右子树高了，模拟插入右儿子
            if (((AVLTreeNode) node.right).heightDif > 0) {
                // 模拟的是插入右儿子的左子树
                AVLTreeNode.rightRound((AVLTreeNode) node.right);
                AVLTreeNode.leftRound(node);
            } else {
                // 模拟的是插入右儿子的右子树
                AVLTreeNode.leftRound(node);
            }
        }
    }

    /**
     * 判断一个节点是否是另一个节点的子孙节点
     * @param node1 子孙节点
     * @param node2 祖先节点
     * @return  如果是返回true，反之返回false
     */
    private boolean isChild(AVLTreeNode node1, AVLTreeNode node2) {
        AVLTreeNode p = node1;
        while (p.parent != null) {
            if (p == node2) {
                return true;
            }
            p = p.parent;
        }
        return false;
    }

    /**
     * 维持AVL树平衡
     * @param parent    不平衡节点
     * @param newNode   插入节点
     */
    private void balance(AVLTreeNode parent, AVLTreeNode newNode) {
        if (parent.left != null && parent.left.left != null && isChild(newNode, (AVLTreeNode) parent.left.left)){
            // 插入节点在不平衡节点的左儿子的左子树上(LL)
            AVLTreeNode.rightRound(parent);
        } else if (parent.right != null && parent.right.right != null && isChild(newNode, (AVLTreeNode) parent.right.right)){
            // 插入节点在不平衡节点的右儿子的右子树上(RR)
            AVLTreeNode.leftRound(parent);
        } else if (parent.left != null && parent.left.right != null && isChild(newNode, (AVLTreeNode) parent.left.right)){
            // 插入节点在不平衡节点的左儿子的右子树上(LR)
            AVLTreeNode.leftRound((AVLTreeNode) parent.left);
            AVLTreeNode.rightRound(parent);
        } else {
            // 插入节点在不平衡节点的右儿子的左子树上(RL)
            AVLTreeNode.rightRound((AVLTreeNode) parent.right);
            AVLTreeNode.leftRound(parent);
        }

        // 如果parent节点原先是根节点则需要更换根节点
        if (this.root == parent) {
            this.root = parent.parent;
        }
    }


    public static void main(String[] args) {
        AVLTree tree = new AVLTree();
        for (int i = 9; i >= 0; --i) {
            tree.insert(i);
        }
        AVLTreeNode.prePrintTree(tree.root);
        boolean success = tree.remove(6);
        System.out.println("success = " + success);
        AVLTreeNode.prePrintTree(tree.root);
        success = tree.remove(11);
        System.out.println("success = " + success);
        AVLTreeNode.prePrintTree(tree.root);
        boolean contain = tree.contains(6);
        System.out.println("contain = " + contain);
        contain = tree.contains(5);
        System.out.println("contain = " + contain);
    }
}
